
# client.py - Demo client for Floral Art API
import os
import sys
import json
import base64
import requests
from typing import Dict, Any, Optional
import argparse
import time
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
API_URL = os.getenv("API_URL", "http://localhost:8000")


class FloralArtClient:
    def __init__(self, api_url: str = API_URL):
        self.api_url = api_url

    def query(self, query_text: str, generate_image: bool = False) -> Dict[str, Any]:
        """Send a query to the API"""
        endpoint = f"{self.api_url}/query"
        payload = {
            "query": query_text,
            "generate_image": generate_image
        }

        try:
            response = requests.post(endpoint, json=payload)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error sending query: {str(e)}")
            return {"error": str(e)}

    def upload_knowledge(self, file_path: str) -> Dict[str, Any]:
        """Upload a document to the knowledge base"""
        if not os.path.exists(file_path):
            return {"error": f"File not found: {file_path}"}

        endpoint = f"{self.api_url}/upload-knowledge"

        try:
            with open(file_path, "rb") as f:
                files = {"file": (os.path.basename(file_path), f)}
                response = requests.post(endpoint, files=files)
                response.raise_for_status()
                return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error uploading file: {str(e)}")
            return {"error": str(e)}

    def transcribe_audio(self, audio_file: str) -> Dict[str, Any]:
        """Transcribe audio file to text"""
        if not os.path.exists(audio_file):
            return {"error": f"Audio file not found: {audio_file}"}

        endpoint = f"{self.api_url}/transcribe-audio"

        try:
            with open(audio_file, "rb") as f:
                audio_data = f.read()
                audio_base64 = base64.b64encode(audio_data).decode("utf-8")

            payload = {"audio_base64": audio_base64}
            response = requests.post(endpoint, json=payload)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error transcribing audio: {str(e)}")
            return {"error": str(e)}

    def generate_image(self, prompt: str, style: str = "realistic") -> Dict[str, Any]:
        """Generate a floral arrangement image"""
        endpoint = f"{self.api_url}/generate-image"
        payload = {
            "prompt": prompt,
            "style": style
        }

        try:
            response = requests.post(endpoint, json=payload)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error generating image: {str(e)}")
            return {"error": str(e)}
    def train_model(self) -> Dict[str, Any]:
        """Start training the local model"""
        endpoint = f"{self.api_url}/train-local-model"
        
        try:
            response = requests.post(endpoint)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error training model: {str(e)}")
            return {"error": str(e)}
            
    def get_system_status(self) -> Dict[str, Any]:
        """Get system status"""
        endpoint = f"{self.api_url}/system-status"
        
        try:
            response = requests.get(endpoint)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error getting system status: {str(e)}")
            return {"error": str(e)}

def main():
    """Main function to demonstrate client usage"""
    parser = argparse.ArgumentParser(description="Floral Art API Client")
    subparsers = parser.add_subparsers(dest="command", help="Command to execute")
    
    # Query command
    query_parser = subparsers.add_parser("query", help="Send a text query")
    query_parser.add_argument("text", help="Query text")
    query_parser.add_argument("--image", action="store_true", help="Generate image")
    
    # Upload command
    upload_parser = subparsers.add_parser("upload", help="Upload document to knowledge base")
    upload_parser.add_argument("file", help="File path to upload")
    
    # Transcribe command
    transcribe_parser = subparsers.add_parser("transcribe", help="Transcribe audio to text")
    transcribe_parser.add_argument("audio", help="Audio file path")
    
    # Generate image command
    image_parser = subparsers.add_parser("image", help="Generate floral image")
    image_parser.add_argument("prompt", help="Image prompt")
    image_parser.add_argument("--style", default="realistic", help="Image style")
    
    # Train model command
    train_parser = subparsers.add_parser("train", help="Train local model")
    
    # Status command
    status_parser = subparsers.add_parser("status", help="Get system status")
    
    # Parse arguments
    args = parser.parse_args()
    
    # Create client
    client = FloralArtClient()
    
    # Process command
    if args.command == "query":
        result = client.query(args.text, args.image)
        print("\nResponse:")
        print(f"Source: {result.get('source', 'unknown')}")
        print(f"Content: {result.get('response', '')}")
        if args.image and "image_url" in result:
            print(f"Image URL: {result['image_url']}")
            
    elif args.command == "upload":
        result = client.upload_knowledge(args.file)
        print(f"\nUpload result: {result.get('message', result)}")
        
    elif args.command == "transcribe":
        result = client.transcribe_audio(args.audio)
        print("\nTranscription:")
        print(result.get("transcription", "Error: No transcription returned"))
        
    elif args.command == "image":
        result = client.generate_image(args.prompt, args.style)
        print("\nImage Generation:")
        print(f"Enhanced prompt: {result.get('prompt', '')}")
        print(f"Image URL: {result.get('image_url', 'None')}")
        
    elif args.command == "train":
        result = client.train_model()
        print(f"\nTraining result: {result.get('message', result)}")
        
    elif args.command == "status":
        result = client.get_system_status()
        print("\nSystem Status:")
        print(f"Knowledge Base: {result.get('knowledge_base', {})}")
        print(f"Local Model Available: {result.get('local_model', {}).get('available', False)}")
        print(f"API Status: {result.get('api_status', 'unknown')}")
        
    else:
        parser.print_help()

if __name__ == "__main__":
    main()
